import math
import typing

from coin.proto.coin_query_pb2 import (AccountBalance, AccountPosition, CurrencyBalance)


def clear_zero_position(account_position: AccountPosition):
  assert type(account_position) == AccountPosition, type(account_position)
  cleared_position = AccountPosition()
  cleared_position.CopyFrom(account_position)
  cleared_position.ClearField('each_position')
  cleared_position.ClearField('each_margin_usage')
  for position in account_position.each_position:
    try:
      if position.net_position != 0 or position.net_position_contract != 0 or \
         position.long_position != 0 or position.short_position != 0 or \
         position.realized_pnl != 0 or position.unrealized_pnl != 0:
        cleared_position.each_position.add().CopyFrom(position)
    except Exception:
      continue
  for margin in account_position.each_margin_usage:
    try:
      if margin.total_margin != 0:
        cleared_position.each_margin_usage.add().CopyFrom(margin)
    except Exception:
      continue
  return cleared_position


def clear_zero_balance(account_balance: AccountBalance):
  assert type(account_balance) == AccountBalance, type(account_balance)
  cleared_balance = AccountBalance()
  cleared_balance.CopyFrom(account_balance)
  cleared_balance.ClearField('each_balance')
  for balance in account_balance.each_balance:
    try:
      if not math.isclose(balance.total, 0):
        cleared_balance.each_balance.add().CopyFrom(balance)
    except Exception:
      continue
  return cleared_balance


def convert_to_balance_dict(account_balance: AccountBalance,
                            nonzero_only: bool=False,) -> typing.Dict[str, float]:
  assert type(account_balance) == AccountBalance, type(account_balance)
  if nonzero_only:
    return {
        balance.currency: balance.total
        for balance in account_balance.each_balance
        if not math.isclose(balance.total, 0, abs_tol=1e-5)
    }
  else:
    return {balance.currency: balance.total for balance in account_balance.each_balance}


def sum_balances(account_balance_list: typing.List[AccountBalance]) -> AccountBalance:
  exchange_list = []
  market_type_list = []
  currency_balances = {}
  for account_balance in account_balance_list:
    assert type(account_balance) == AccountBalance, type(account_balance)
    exchange_list.append(account_balance.exchange)
    market_type_list.append(account_balance.market_type)
    acct_ccy_set = set()
    for balance in account_balance.each_balance:
      assert balance.currency not in acct_ccy_set, f'{account_balance.market_type}_{account_balance.exchange}_{balance.currency}'
      acct_ccy_set.add(balance.currency)
      if balance.currency in currency_balances:
        currency_balances[balance.currency].total += balance.total
        currency_balances[balance.currency].available += balance.available
        currency_balances[balance.currency].hold += balance.hold
      else:
        currency_balances[balance.currency] = CurrencyBalance(
            currency=balance.currency,
            total=balance.total,
            available=balance.available,
            hold=balance.hold,
        )
  exchange = exchange_list[0] if len(set(exchange_list)) == 1 else None
  market_type = market_type_list[0] if len(set(market_type_list)) == 1 else None
  return AccountBalance(
      exchange=exchange,
      market_type=market_type,
      each_balance=currency_balances.values(),
  )


def gen_balance_with_reserve(reserve_config):
  balance_list = []
  for currency, balance in reserve_config.items():
    currency_balance = CurrencyBalance(
        currency=currency,
        total=balance,
        available=balance,
        hold=0.0,
    )
    balance_list.append(currency_balance)
  return AccountBalance(each_balance=balance_list)
